package MLDA_Process;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import Commons.IOUtil;


public class MixtureOfUnigrams implements Runnable {
    public int K = 100;

    public double alpha = 0.1;
    public double beta = 0.01;

    public int niters = 2000;

    public String inputPath = "";
    public String outputPath = "";

    public int[] zAssigns;

    public int nk[];
    public int nkw[][];
    public int nkwSum[];

    public int M;
    public int V;
    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();

    public MixtureOfUnigrams(int k, int niters, String inputPath, String outputPath) {
        this.K = k;
        this.niters = niters;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }


    public void loadTxts(String txtPath) {
        String line;
        try {
            BufferedReader reader = new BufferedReader(new FileReader(txtPath));
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();
                String[] words = line.trim().split("\\s+");
                for (String word : words) {
                    if (!w2i.containsKey(word)) {
                        w2i.put(word, w2i.size());
                        i2w.put(w2i.get(word), word);
                    }
                    doc.add(w2i.get(word));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        M = docs.size();
        V = w2i.size();

        System.out.println(M);
        System.out.println(V);

        return;
    }

    public void initMU() {
        nk = new int[K];
        nkw = new int[K][V];
        nkwSum = new int[K];

        zAssigns = new int[M];

        for (int m = 0; m != M; m++) {
            int z = (int) Math.floor(Math.random() * K);

            int N = docs.get(m).size();
            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                nkw[z][w]++;
                nkwSum[z]++;
            }
            nk[z]++;
            zAssigns[m] = z;
        }

        return;
    }

    public void sampleZ(int m) {
        HashMap<Integer, Integer> w2c = new HashMap<Integer, Integer>();
        int z = zAssigns[m];
        int N = docs.get(m).size();
        for (int n = 0; n != N; n++) {
            int w = docs.get(m).get(n);
            if (!w2c.containsKey(w)) {
                w2c.put(w, 1);
            } else {
                w2c.put(w, w2c.get(w) + 1);
            }
            nkw[z][w]--;
            nkwSum[z]--;
        }
        nk[z]--;

        double VBeta = V * beta;
        double[] pTable = new double[K];

        for (int k = 0; k != K; k++) {
            List<Double> devidents = new ArrayList<Double>();

            for (int n = 0; n != N; n++) {
                devidents.add(VBeta + nkwSum[k] + n);
            }

            double expectTM = 1.0;
            int index = 0;
            for (int w : w2c.keySet()) {
                int c = w2c.get(w);
                for (int i = 0; i != c; i++) {
                    expectTM *= (nkw[k][w] + beta + i) / devidents.get(index++);
                }
            }

            pTable[k] = (nk[k] + alpha) / (M + K * alpha) * expectTM;
        }

        for (int k = 1; k != K; k++) {
            pTable[k] += pTable[k - 1];
        }

        double r = Math.random() * pTable[K - 1];

        for (int k = 0; k != K; k++) {
            if (pTable[k] > r) {
                z = k;
                break;
            }
        }

        for (int n = 0; n != N; n++) {
            int w = docs.get(m).get(n);
            nkw[z][w]++;
            nkwSum[z]++;
        }
        nk[z]++;

        zAssigns[m] = z;
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("MU Iteration: " + iter + " ...");
            for (int m = 0; m != M; m++) {
                sampleZ(m);
            }
        }
        return;
    }

    public double[][] computePhi() {
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nkw[k][v] + beta) / (nkwSum[k] + V * beta);
            }
        }
        return phi;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            res.add(pairs);
        }

        return res;
    }

    public void print_topics(int top_n) {
        double phi[][] = computePhi();
        ArrayList<List<Entry<String, Double>>> pairs_list = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            System.out.println("Topic " + k + ":");
            for (int i = 0; i != top_n; i++) {
                System.out.println(pairs_list.get(k).get(i).getKey() + " "
                        + pairs_list.get(k).get(i).getValue());
            }
        }
    }

    public void print_topics(String path, int top_n) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path);
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
        writer.close();
    }

    public void saveZAssigns(String filename) {
        BufferedWriter writer = IOUtil.getWriter(filename);
        try {
            for (int z : zAssigns) {
                writer.append(z + "\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return;
    }

    public double[][] computeTheta() {
        double[][] theta = new double[M][K];
        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                if (zAssigns[m] == k) {
                    theta[m][k] = ((double) docs.get(m).size() + alpha) / ((double) docs.get(m).size() + K * alpha);
                } else {
                    theta[m][k] = alpha / K * alpha;
                }
            }
        }
        return theta;
    }

    public void saveTheta(String path) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = this.computeTheta();

        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                writer.append(theta[m][k] + " ");
            }
            writer.newLine();
        }
        writer.close();
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();
        int K = phi.length;
        assert K > 0;
        int V = phi[0].length;

        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++) {
                    writer.append(phi[k][v] + " ");
                }
                writer.append("\n");
            }

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return;
    }

    public void printModel() {
        System.out.println("niters:" + this.niters + "\tbeta:" + this.beta + "\talpha:" + this.alpha + "\tk:" + this.K + "\tinput:" + this.inputPath + "\toutput:" + this.outputPath);
    }

    public void run() {
        printModel();
        this.loadTxts(this.inputPath);
        this.initMU();
        this.estimate();
        this.savePhi(this.outputPath + "model-final.phi");
        this.saveWordmap(this.outputPath + "wordmap.txt");
        try {
            this.print_topics(this.outputPath + "model-final.twords", 20);
            this.saveTheta(this.outputPath + "model-final.theta");
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        saveZAssigns(this.outputPath + "assign.txt");
    }

    public static void UnigramText(int k, int iter, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath();
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + "/unigram/" + i)).mkdirs();
            service.submit(new MixtureOfUnigrams(k, iter, path, parent_path + "/unigram/" + i + "/"));

        }
        service.shutdown();
    }
}
